﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.IO;

namespace Kokomo.Mirror.Ia32
{
	public class IaDisassembler
	{
		static IaDisassembler()
		{
			Disassembler32 = new IaDisassembler(OperandSize.Dword, Ia32OpCodeLookup.Instance);
			Disassembler64 = new IaDisassembler(OperandSize.Qword, Ia32eOpCodeLookup.Instance);
		}

		public static IaDisassembler Disassembler32 { get; private set; }
		public static IaDisassembler Disassembler64 { get; private set; }

		internal IaDisassembler(
			OperandSize defaultRegisterSize,
			IaOpCodeLookup opCodeLookup
			)
		{
			if (opCodeLookup == null) throw new ArgumentNullException("opCodeLookup");

			this.OpCodeLookup = opCodeLookup;
			this.DefaultRegisterSize = defaultRegisterSize;
		}

		public IaOpCodeLookup OpCodeLookup { get; private set; }
		public OperandSize DefaultRegisterSize { get; private set; }

		private static long NextOffsetAfter(
			List<long> offsets,
			long offset
			)
		{
			return offsets.Where(o => o > offset).OrderBy(o => o).FirstOrDefault();
		}
		public Ia32Disassembly Disassemble(
			Stream stream,
			long baseAddress
			)
		{
			if (stream == null) throw new ArgumentNullException("stream");
			if (stream == null) throw new ArgumentNullException("stream");

			List<NativeDependency> dependencies = new List<NativeDependency>();

			bool done = false;
			List<Ia32DisassemblyBlock> blocks = new List<Ia32DisassemblyBlock>();

			long streamBaseOffset = baseAddress - stream.Position;

			List<long> jumpOffsets = new List<long>();
			Stack<long> blockOffsets = new Stack<long>();
			blockOffsets.Push(baseAddress);
			while (blockOffsets.Count > 0)
			{
				long blockBaseAddress = blockOffsets.Pop();
				if (jumpOffsets.Contains(blockBaseAddress)) continue;

				long? nextAddress = null;

				bool blockAlreadyDisassembled = false;
				foreach (var block in blocks)
				{
					if (block.BaseAddress == blockBaseAddress)
					{
						blockAlreadyDisassembled = true;
						break;
					}
					if (block.ContainsAddress(blockBaseAddress))
					{
						Ia32DisassemblyBlock before, after;
						SplitBlock(block, blockBaseAddress, out before, out after);

						int blockPosition = blocks.IndexOf(block);
						blocks.Insert(blockPosition, before);
						blocks.Insert(blockPosition + 1, after);
						blocks.RemoveAt(blockPosition + 2);
						blockAlreadyDisassembled = true;
						break;
					}

					if (block.BaseAddress > blockBaseAddress)
					{
						if (nextAddress.HasValue)
							nextAddress = Math.Min(nextAddress.Value, block.BaseAddress);
						else
							nextAddress = block.BaseAddress;
					}
				}
				if (blockAlreadyDisassembled) continue;

				int? lengthLimit = nextAddress.HasValue
					? (int?)(nextAddress.Value - blockBaseAddress)
					: null
					;
				stream.Position = blockBaseAddress - streamBaseOffset;
				Ia32DisassemblyBlock newBlock = this.DisassembleBlock(
					stream,
					blockBaseAddress,
					dependencies,
					lengthLimit
					);
				blocks.Add(newBlock);

				jumpOffsets.Add(blockBaseAddress);

				if (newBlock.HasFallthrough)
				{
					blockOffsets.Push(newBlock.BaseAddress + newBlock.Length);
				}
				if (newBlock.HasFinalJump && newBlock.FinalJumpAddress != 0)
				{
					blockOffsets.Push(newBlock.FinalJumpAddress);
				}
			}

			Ia32Disassembly disassembly = new Ia32Disassembly(
				baseAddress,
				0,
				blocks,
				dependencies
				);

			return disassembly;
		}
		private static void SplitBlock(
			Ia32DisassemblyBlock block,
			long address,
			out Ia32DisassemblyBlock before,
			out Ia32DisassemblyBlock after
			)
		{
			if (block == null) throw new ArgumentNullException("block");
			if (!block.ContainsAddress(address)) throw new ArgumentException("block does not contain address");

			List<Ia32Instruction> beforeInstructions = new List<Ia32Instruction>();
			List<Ia32Instruction> afterInstructions = new List<Ia32Instruction>();
			foreach (var instruction in block.Instructions)
			{
				if (instruction.Offset < address)
					beforeInstructions.Add(instruction);
				else
					afterInstructions.Add(instruction);
			}

			before = new Ia32DisassemblyBlock(
				block.BaseAddress,
				(int)(address - block.BaseAddress),
				beforeInstructions
				);
			after = new Ia32DisassemblyBlock(
				address,
				block.Length - before.Length,
				afterInstructions
				);
		}
		private static bool EndsBlock(FlowControl flow)
		{
			return
				(flow != FlowControl.Normal)
				&& (flow != FlowControl.Call)
				;
		}
		private Ia32DisassemblyBlock DisassembleBlock(
			Stream stream,
			long blockBaseAddress,
			IList<NativeDependency> dependencies,
			int? blockLengthLimit = null
			)
		{
			long streamStartPosition = stream.Position;
			long streamBaseAddress = blockBaseAddress - stream.Position;
			long? streamEndPositionLimit = (!blockLengthLimit.HasValue) ? (long?)null : stream.Position + blockLengthLimit.Value;
			IaReader reader = new IaReader(stream, streamBaseAddress, this.DefaultRegisterSize, this.OpCodeLookup);
			List<Ia32Instruction> instructions = new List<Ia32Instruction>();

			bool done = false;

			while (
				!done
				&& (!streamEndPositionLimit.HasValue || (stream.Position < streamEndPositionLimit))
				&& reader.MoveNext()
				)
			{
				Ia32Instruction instruction = reader.CurrentInstruction;
				instructions.Add(instruction);

				FlowControl flow = instruction.OpCode.Mnemonic.FlowControl;

				done = EndsBlock(flow);
				if (flow != FlowControl.Normal)
				{
					JumpTarget branch = null;
					if (instruction.Operands.Length == 1)
					{
						branch = instruction.Operands[0] as JumpTarget;
					}

					if (branch != null)
					{
						if (flow == FlowControl.Call)
						{
							NativeDependency dependency = new NativeProcedureDependency(instruction.Offset, branch.TargetOffset);
							if (dependencies != null) dependencies.Add(dependency);
						}
					}
				}
			}

			int blockLength = (int)(stream.Position - streamStartPosition);
			Ia32DisassemblyBlock block = new Ia32DisassemblyBlock(
				blockBaseAddress,
				blockLength,
				instructions
				);

			return block;
		}
	}
}
